
import tempfile
from files import FileTool
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import numpy as np
import logging
import copy
from transformer import to_npy, DataTransformer
import tempfile
import argparse
import pydicom
parser = argparse.ArgumentParser(description='format convert: dicom zip to nii')
parser.add_argument('-i','--input_dir', help='Path to input image',required=False)
parser.add_argument('-o','--output_dir',help='Path th output image',required=False)
args = parser.parse_args()
FORMAT = '%(asctime)s - %(levelname)s: %(message)s'
logging.basicConfig(format=FORMAT,datefmt="%Y-%m-%d %H:%M:%S",level='INFO')

class MaskConverter:
    def __init__(self, filepath) -> None:
        """filepath 模型结果文件路径"""
        # self.files_list = None
        base_dir = Path(__file__).parent
        self.filepath = Path(filepath)

    def extra_cond(self, p):
        return p.stem == 'payload'

    def get_volumes_from_series(self,series_files):
        instance_infos = []
        for file_name in series_files:
            ds = pydicom.dcmread(
                file_name, 
                force=True, 
                stop_before_pixels=True
            )
            if not hasattr(ds, 'InstanceNumber'):
                print(f'warning: InstanceNumber not exist {file_name}')
                continue
            if not hasattr(ds, 'ImagePositionPatient'):
                print(f'warning: ImagePositionPatient not exist {file_name}')
                continue
            instance_number = ds.get('InstanceNumber')
            image_position_patient = list(ds.get("ImagePositionPatient"))
            info = {
                    "instance_number": instance_number,
                    "image_position_patient": np.array(image_position_patient),
                    "file_name": file_name
                }
            instance_infos.append(info)

        def func(f):
            return f.get("instance_number", -1)
        sorted_instance_infos = sorted(instance_infos, key=func, reverse=False)

        v_ref = sorted_instance_infos[0].get('image_position_patient') - sorted_instance_infos[1].get('image_position_patient')
        volumes_files = []

        vol_files = [sorted_instance_infos[0], sorted_instance_infos[1]]
        for i in range(2, len(sorted_instance_infos)):
            v = sorted_instance_infos[i-1].get('image_position_patient') - sorted_instance_infos[i].get('image_position_patient')
            if np.dot(v, v_ref) > 0:
                vol_files.append(sorted_instance_infos[i])
            else:
                volumes_files.append(copy.deepcopy(vol_files))
                vol_files.clear()
                vol_files.append(sorted_instance_infos[i])

        volumes_files.append(copy.deepcopy(vol_files))
        return volumes_files

    def iter_files(self, p, results, extra_cond=None):
        if p.is_file():
            if extra_cond and self.extra_cond(p) or not extra_cond:
                if p.suffix == '.zip':
                    results.append(p.resolve())
                # elif p.suffix == '.pkl':
                #     content = FileTool.read_pkl(p.resolve().as_posix())
                else:
                    return

        elif p.is_dir():
            sub_res = []
            for f in p.iterdir():
                self.iter_files(f, sub_res, extra_cond)
            if sub_res:
                self.generate_nii(sub_res)
            for _ in range(len(sub_res)):
                sub_res.pop()
    
    def generate_nii(self, zip_list):
        # for files in iter_files:
        res=None
        logging.info(f"get data: {zip_list[0].parent.name}")
        nii_path = zip_list[0].parent / "nii_path"
        FileTool.make_directory(nii_path)
        nii_index = 0
        for i, file in enumerate(zip_list, start=1):
            logging.info(f"-----get label:{file.stem}------")
            with tempfile.TemporaryDirectory() as temp_obj:
                FileTool.unzip_file(file.as_posix(), temp_obj)
                file_list = list(Path(temp_obj).iterdir())
                file_list = self.get_volumes_from_series(file_list)
                # file_list = [i.resolve().as_posix() for i in file_list]
                file_list = [str(file.get('file_name')) for file in file_list[0]]
                sitk_obj,arr = to_npy(file_list)
                arr[arr > 0] = i
                if res is not None and res.any():
                    if np.logical_and(res>0, arr>0).any():
                        logging.info("有重叠区域，生成单独nii")
                        sitk_obj = DataTransformer.NpyToSitk(arr, sitk_obj)
                        DataTransformer.SitkToNii(sitk_obj, str(nii_path / f"{str(nii_index)}.nii.gz"))
                        nii_index+=1
                    else:
                        res = np.where(arr>0, arr, res)
                    # res[arr>0]=arr
                else:
                    res = arr
                logging.info(f"{i}, {[i for i in range(32) if(res==i).any()]}")
        if res is not None:
            sitk_obj = DataTransformer.NpyToSitk(res, sitk_obj)
            DataTransformer.SitkToNii(sitk_obj, str(nii_path / f"{str(nii_index)}.nii.gz"))
            logging.info('generate nii file finished')
    def main(self, outdir=None):
        """outdir 注释文件输出路径，不提供时默认文件生成路径为模型所在位置"""
        iter_files = []
        self.iter_files(self.filepath, iter_files)
        return 


if __name__ == '__main__':
    input_dir = '/Users/hongwei/workspace/develop_tools/dcm2nii_model/data_repo/945458'
    output_dir = ''
    if args.input_dir:
        input_dir = args.input_dir
    if args.output_dir:
        output_dir = args.output_dir
    
    generator = MaskConverter(input_dir)
    generator.main(output_dir)
